########### Packages
import pandas as pd
import numpy as np
from scipy import stats
import researchpy as rp
import pystan
import pystan.experimental
import stan_utility
### Create Data Dictionaries ###############3
reg_data2_pg=read_csv('full_replication_data.csv')
y_pg=reg_data2_pg['pct_female_cand']
y_logit_pg=reg_data2_pg['new_quota']
x2_pg=reg_data2_pg[[ 'PerUrb' , 'PerUkrspk', 'unem2014', 'pct_female_dep_10',  'm_2015', 'ln_days_old', 'pct_districts']]
x2_pg=x2_pg.apply(lambda x2_pg: x2_pg-x2_pg.mean())
left_p_pg=reg_data2_pg[['left']].astype(int)
x2_pg.insert(3, 'left', value=left_p_pg)
en_p_pg=reg_data2_pg[['en_party']].astype(int)
x2_pg.insert(3,'en_party', value=en_p_pg)
r14_pg=reg_data2_pg[['rada_14_1pct']].astype(int)
#x2_pg.insert(3,'rada_14_1pct', value=r14_pg)
pct_10=reg_data2_pg[['pct_total_2010']]
x2_pg.insert(9, 'pct_total_2010', value=pct_10)
r_pv_14=reg_data2_pg[['r_pv_14']]
x2_pg.insert(10, 'r_pv_14', value=r_pv_14)
df_c_pg=rp.summary_cat(reg_data2_pg['region_3'])
df_c_pg.loc[:,'city']=(df_c_pg.index)+1
df_c_pg=df_c_pg.rename(columns={'Outcome':'region_3'})
reg_data2_pg=pd.merge(reg_data2_pg, df_c_pg, on=['region_3'], how='left')
city_pg=reg_data2_pg['city']
party_pg=reg_data2_pg['all_party'].astype(int)
k2_pg=len(x2_pg.columns)
p_pg=reg_data2_pg['all_party'].max()
c_pg=reg_data2_pg['city'].max()
dict_m2_pg={'N':len(reg_data2_pg.index), 'K':k2_pg, 'C':c_pg, 'P':p_pg, 'y':y_pg, 'x':x2_pg, 'city':city_pg, 'party':party_pg}
dict_logit2_pg={'N':len(reg_data2_pg.index), 'K':k2_pg, 'C':c_pg, 'P':p_pg, 'y':y_logit_pg, 'x':x2_pg, 'city':city_pg, 'party':party_pg}

######### Model 1 #####################
hlm_model_1=pystan.StanModel(file="model_1.stan")

def init_values():
     return  dict(sigma_y=stats.halfcauchy.rvs(0,2.5), sigma_c=stats.halfnorm.rvs(0, 1),
                  alpha=stats.norm.rvs(0,1), sigma_p=stats.halfnorm.rvs(0, 1), nu=stats.halfnorm.rvs(2,1),
                 beta=np.random.normal(0,1,size=(k2_pg)))

fit_hlm_model_1=hlm_model_1.sampling(data=dict_m2_pg, warmup=5000, iter=10000, chains=4,
                                   seed=8915356, n_jobs=4, thin=10,
                                   control=dict(max_treedepth=18, adapt_delta=0.9),
                                  init=init_values)

########## Model 2 ##########################
logit_model_2=pystan.StanModel(file="logit_model_2.stan")

def init_values():
     return  dict (sigma_c=stats.halfnorm.rvs(0, 1), alpha=stats.norm.rvs(0,1),
                   sigma_p=stats.halfnorm.rvs(0, 1), nu=stats.halfnorm.rvs(5,1),
                  beta=np.random.normal(0,1,size=(k2_pg)))
fit_mle_model_2=logit_model_2.sampling(data=dict_logit2_pg, warmup=5000, iter=10000,
                                     chains=4, seed=8915356, n_jobs=4, thin=10,
                                     control=dict(max_treedepth=18, adapt_delta=0.99),
                                    init=init_values)
